# Import required libraries
import marimo as mo
import altair as alt
import polars as pl
import numpy as np
import requests
import folium
from typing import Optional

def create_cog_example_data():
    """Create 5 example demand points for CoG demonstration."""
    return pl.DataFrame({
        'id': [1, 2, 3, 4, 5],
        'x': [10, 50, 20, 90, 80],
        'y': [20, 20, 80, 90, 5],
        'demand': [50, 200, 300, 80, 400]
    })

def create_interactive_cog_chart():
    """
    Create an interactive chart where points can be clicked to select/unselect.
    The CoG is calculated and displayed dynamically based on selection.
    """
    points = create_cog_example_data()#.to_pandas()

    # Create click selection (multi-select with toggle, start empty)
    # toggle='true' (string) makes every click toggle without needing shift key
    # empty=False means no points are selected initially
    click = alt.selection_point(fields=['id'], toggle='true', clear=False, empty=False)

    # Demand points
    points_chart = alt.Chart(points).mark_circle().encode(
        x=alt.X('x:Q', scale=alt.Scale(domain=[0, 100]), title='X Coordinate'),
        y=alt.Y('y:Q', scale=alt.Scale(domain=[0, 100]), title='Y Coordinate'),
        size=alt.Size('demand:Q', scale=alt.Scale(range=[200, 1000]), legend=None),
        color=alt.condition(
            click,
            alt.value('steelblue'),
            alt.value('lightgray')
        ),
        opacity=alt.condition(
            click,
            alt.value(0.8),
            alt.value(0.3)
        ),
        tooltip=[
            alt.Tooltip('id:O', title='Point ID'),
            alt.Tooltip('x:Q', title='X'),
            alt.Tooltip('y:Q', title='Y'),
            alt.Tooltip('demand:Q', title='Demand')
        ]
    ).add_params(click)

    # Point labels
    labels = alt.Chart(points).mark_text(dy=-15, fontSize=12, fontWeight='bold').encode(
        x='x:Q',
        y='y:Q',
        text='id:O',
        color=alt.condition(
            click,
            alt.value('black'),
            alt.value('gray')
        )
    ).add_params(click)

    # Calculate CoG from selected points using Altair transforms
    cog_chart = alt.Chart(points).transform_filter(
        click
    ).transform_calculate(
        xq='datum.x * datum.demand',
        yq='datum.y * datum.demand'
    ).transform_aggregate(
        sum_xq='sum(xq)',
        sum_yq='sum(yq)',
        sum_q='sum(demand)',
        groupby=[]
    ).transform_calculate(
        cog_x='datum.sum_xq / datum.sum_q',
        cog_y='datum.sum_yq / datum.sum_q'
    ).mark_point(
        shape='cross',
        size=500,
        color='red',
        strokeWidth=4
    ).encode(
        x='cog_x:Q',
        y='cog_y:Q',
        tooltip=[
            alt.Tooltip('cog_x:Q', title='CoG X', format='.1f'),
            alt.Tooltip('cog_y:Q', title='CoG Y', format='.1f')
        ]
    )

    # CoG label
    cog_label = alt.Chart(points).transform_filter(
        click
    ).transform_calculate(
        xq='datum.x * datum.demand',
        yq='datum.y * datum.demand'
    ).transform_aggregate(
        sum_xq='sum(xq)',
        sum_yq='sum(yq)',
        sum_q='sum(demand)',
        groupby=[]
    ).transform_calculate(
        cog_x='datum.sum_xq / datum.sum_q',
        cog_y='datum.sum_yq / datum.sum_q'
    ).mark_text(
        dy=-15,
        fontSize=14,
        fontWeight='bold',
        color='red'
    ).encode(
        x='cog_x:Q',
        y='cog_y:Q'
    )

    # Lines from CoG to selected points
    lines_base = alt.Chart(points).transform_filter(
        click
    ).transform_calculate(
        xq='datum.x * datum.demand',
        yq='datum.y * datum.demand'
    ).transform_joinaggregate(
        sum_xq='sum(xq)',
        sum_yq='sum(yq)',
        sum_q='sum(demand)'
    ).transform_calculate(
        cog_x='datum.sum_xq / datum.sum_q',
        cog_y='datum.sum_yq / datum.sum_q'
    )

    # Draw lines
    lines = lines_base.mark_line(
        color='red',
        strokeWidth=2,
        opacity=0.4,
        strokeDash=[5, 5]
    ).encode(
        x='x:Q',
        y='y:Q',
        x2='cog_x:Q',
        y2='cog_y:Q'
    )

    # Demand labels on lines (at midpoint)
    line_labels = lines_base.transform_calculate(
        mid_x='(datum.x + datum.cog_x) / 2',
        mid_y='(datum.y + datum.cog_y) / 2'
    ).mark_text(
        fontSize=11,
        fontWeight='bold',
        color='red',
        dx=0,
        dy=-5
    ).encode(
        x='mid_x:Q',
        y='mid_y:Q',
        text=alt.Text('demand:Q', format='d')
    )

    return (points_chart + labels + lines + line_labels + cog_chart + cog_label).properties(
        width=500,
        height=500,
        title='Click points to select/unselect for CoG calculation'
    )

_cog_example_plot = create_interactive_cog_chart()

_cog_example_plot